#!python3

import sys
import hashlib
from base64 import b64encode, b64decode
from collections import deque
import re

prooffile = "proof.txt"

MAXHEIGHT = 20

class MerkleProof:
    def __init__(self, leaf, pos, path):
        self.leaf = leaf
        self.pos  = pos
        self.path = path


def hash_leaf(leaf):
    sha256 = hashlib.sha256()
    sha256.update(b"leaf:")
    sha256.update(leaf)
    return sha256.digest()

def hash_internal_node(left, right):
    sha256 = hashlib.sha256()
    sha256.update(b"node:")
    sha256.update(left)
    sha256.update(right)
    return sha256.digest()


def compute_merkle_root_from_proof(proof: MerkleProof):
    pos = proof.pos
    path = deque(proof.path)
    root = hash_leaf(proof.leaf)
    while path:
        if pos % 2 == 0:
            left, right = root, path.popleft()
        else:
            left, right = path.popleft(), root
        root = hash_internal_node(left, right)
        pos >>= 1
    return root


def read_proof(filename):
    fp = open(filename, "r")
    pos  = int(re.search('(\d*)$', fp.readline()).group(1))
    leaf = re.search('\"(.*)\"',fp.readline()).group(1).encode()
    fp.readline()
    path = fp.readlines()
    for i in range(len(path)):
        path[i] = b64decode((path[i])[2:])
    fp.close()
    return MerkleProof(leaf=leaf, pos=pos, path=path)


if __name__ == "__main__":
    root = b64decode("1qIbsvuF6FrhNjMD4p06srUye6G4FfFINDDkNfKUpTs=")
    print('\nHave hardcoded root of committed Merkle tree.')

    proof = read_proof(prooffile)

    height = len(proof.path)
    assert height < MAXHEIGHT, "Proof is too long"

    computedRoot = compute_merkle_root_from_proof(proof)
    assert root == computedRoot, "Verify failed"

    print('I verified the Merkle proof: leaf #{} in the committed tree is "{}".\n'.format(proof.pos,proof.leaf.decode("utf-8")))
    sys.exit(0)



